import cv2
import numpy as np
import torch
import torchvision.transforms as T
from torchvision.models.detection import maskrcnn_resnet50_fpn
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
from PIL import Image, ImageDraw


def detect_obstacles(image_path, output_path):
    # 加载预训练模型
    model = maskrcnn_resnet50_fpn(pretrained=True)

    # 替换分类器，仅保留实例分割部分
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = MaskRCNNPredictor(in_features, dim_reduced=3, num_classes=2)  # 2 是类别数

    # 加载图像
    image = Image.open(image_path)

    # 图像预处理
    transform = T.Compose([T.ToTensor()])
    image = transform(image)

    # 进行预测
    model.eval()
    with torch.no_grad():
        prediction = model([image])

    # 可视化预测结果
    masks = prediction[0]['masks'].cpu().numpy()
    img = Image.open(image_path)
    draw = ImageDraw.Draw(img)

    for mask in masks:
        mask = mask[0]
        contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for contour in contours:
            area = cv2.contourArea(contour)
            if area < 500:  # 根据实际情况调整遮挡区域的最小面积
                x, y, w, h = cv2.boundingRect(contour)
                draw.rectangle([x, y, x + w, y + h], outline="red", width=3)

    img.save(output_path)


if __name__ == "__main__":
    image_path = 'config/hikcam.jpg'
    image_path2 = 'config/zhedang.jpeg'
    image_path_blur = 'config/blurred_image.jpg'
    image_path_obstruction = 'config/obstruction_image.jpg'

    input_image_path = image_path
    output_image_path = "config/output_image.jpg"
    detect_obstacles(input_image_path, output_image_path)
